from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType


def run_coze_workflow(user_input:str,workflow_id: str, api_token: str, base_url: str = COZE_CN_BASE_URL):
    """
    运行Coze工作流并处理返回的流事件

    参数:
        workflow_id: 工作流ID
        api_token: Coze API访问令牌
        base_url: Coze API基础URL(默认为中国区)

    返回:
        list: 包含所有工作流消息的列表
    """
    # 初始化Coze客户端
    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)

    messages = []  # 用于存储消息的列表

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        """处理工作流事件迭代器"""
        nonlocal messages  # 引用外部变量
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                messages.append(event.message)  # 存储消息
            elif event.event == WorkflowEventType.ERROR:
                raise Exception(f"工作流错误: {event.error}")  # 将错误抛出
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="hey",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 获取用户输入
    # user_input = input("请输入你想搜索的内容: ")

    # 运行工作流并处理结果
    handle_workflow_iterator(
        coze.workflows.runs.stream(
            workflow_id=workflow_id,
            parameters={"input": user_input}
        )
    )

    return messages  # 返回所有消息




# 使用示例
if __name__ == "__main__":
    # 配置参数
    WORKFLOW_ID = "7529426430200987711"
    API_TOKEN = "pat_dCQtfQFawHONOqIyBmRfE1VI8OJLFo0ylgP2UgUr9v0h9mxQk8jL89UoV8fukZWF"

    # 运行工作流
    result=run_coze_workflow("数学",workflow_id=WORKFLOW_ID, api_token=API_TOKEN)
    import json

    content_dict = json.loads(result[0].content)  # 先将字符串转为字典
    print("Result:", content_dict["output"])  # 然后再访问output字段